package com.intelligent.system.common.websocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
public class WebSocketHandler extends TextWebSocketHandler {
    private static final Map<String, WebSocketSession> SESSIONS = new ConcurrentHashMap<>();
    private final ObjectMapper objectMapper = new ObjectMapper();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        String clientId = getClientId(session);
        SESSIONS.put(clientId, session);
        log.info("客户端 {} 的WebSocket连接已建立", clientId);

        // 发送连接成功消息
        sendMessage(session, new WebSocketMessage("system", "server", clientId, "Connected successfully"));
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        try {
            WebSocketMessage webSocketMessage = objectMapper.readValue(message.getPayload(), WebSocketMessage.class);
            String clientId = getClientId(session);
            webSocketMessage.setFrom(clientId);

            // 处理不同类型的消息
            switch (webSocketMessage.getType()) {
                case "private":
                    // 发送私聊消息
                    sendPrivateMessage(webSocketMessage);
                    break;
                case "broadcast":
                    // 发送广播消息
                    broadcastMessage(webSocketMessage);
                    break;
                default:
                    log.warn("未知的消息类型: {}", webSocketMessage.getType());
            }
        } catch (IOException e) {
            log.error("处理消息失败", e);
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) {
        String clientId = getClientId(session);
        log.error("客户端 {} 的WebSocket传输错误", clientId, exception);
        SESSIONS.remove(clientId);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        String clientId = getClientId(session);
        log.info("客户端 {} 的WebSocket连接已关闭，状态: {}", clientId, status);
        SESSIONS.remove(clientId);
    }

    /**
     * 发送私聊消息
     */
    public void sendPrivateMessage(WebSocketMessage message) {
        WebSocketSession targetSession = SESSIONS.get(message.getTo());
        if (targetSession != null && targetSession.isOpen()) {
            sendMessage(targetSession, message);
        } else {
            log.warn("目标客户端未找到或连接已关闭: {}", message.getTo());
        }
    }

    /**
     * 发送广播消息
     */
    public void broadcastMessage(WebSocketMessage message) {
        SESSIONS.forEach((clientId, session) -> {
            if (session.isOpen()) {
                sendMessage(session, message);
            }
        });
    }

    /**
     * 发送消息
     */
    private void sendMessage(WebSocketSession session, WebSocketMessage message) {
        try {
            String messageJson = objectMapper.writeValueAsString(message);
            session.sendMessage(new TextMessage(messageJson));
        } catch (IOException e) {
            log.error("发送消息失败", e);
        }
    }

    /**
     * 获取客户端ID
     */
    private String getClientId(WebSocketSession session) {
        return session.getId();
    }

    /**
     * 获取所有在线客户端
     */
    public Map<String, WebSocketSession> getSessions() {
        return SESSIONS;
    }
}